import sqlite3
from datetime import datetime, timedelta
import pandas as pd
import os
import sys
import akshare as ak
import json
from bk_work_run import is_trade_day

# 数据库配置
DB_FILE = 'stock_data.db'


def read_list_file(file_path):
    """
    加载文件,生成一个列表返回编码和名称,同时可以方便通过编码查询名称
    """
    stock_list = []  # 存储股票信息的列表
    stock_dict = {}  # 存储股票编码到名称的映射字典
    
    try:
        # 尝试不同的编码格式
        encodings = ['utf-8', 'gbk', 'gb2312', 'latin-1']
        file_content = None
        used_encoding = None
        
        for encoding in encodings:
            try:
                with open(file_path, 'r', encoding=encoding) as f:
                    file_content = f.read()  # 一次性读取全部内容
                print(f"成功使用 {encoding} 编码读取文件")
                used_encoding = encoding
                break
            except UnicodeDecodeError:
                continue
        
        # 如果使用的编码不是UTF-8，则转换为UTF-8并保存
        if file_content is not None and used_encoding != 'utf-8':
            print(f"文件使用的是非UTF-8编码 ({used_encoding})，正在转换为UTF-8...")
            # 将内容转换为UTF-8并保存
            with open(file_path, 'w', encoding='utf-8') as f:
                f.write(file_content)
            print(f"文件已成功转换为UTF-8编码并保存")
        
        # 将文件内容分割成行
        file_content = file_content.splitlines() if file_content else []
        
        if file_content is None:
            raise Exception("无法使用常见编码格式读取文件")
        
        for line in file_content:
            # 去除行首尾空白字符
            line = line.strip()
            if line:
                # 分割编码和名称（假设用空格或制表符分隔）
                parts = line.split()
                if len(parts) >= 2:
                    code = parts[0]  # 股票编码
                    # 根据股票代码格式处理规范，去除每行开头的2个字符（如SH/SZ），仅保留6位数字代码
                    if len(code) > 2:
                        code = code[2:]
                    name = ' '.join(parts[1:])  # 股票名称（处理名称中可能包含空格的情况）
                    stock_list.append((code, name))  # 添加到列表
                    stock_dict[code] = name  # 添加到字典，便于通过编码查询名称
                elif len(parts) == 1:
                    # 如果只有一部分，可能只有编码没有名称
                    code = parts[0]
                    stock_list.append((code, ""))  # 名称为空
                    stock_dict[code] = ""  # 添加到字典
    except FileNotFoundError:
        print(f"错误: 文件 {file_path} 不存在")
        return [], {}
    except Exception as e:
        print(f"读取文件 {file_path} 时出错: {e}")
        return [], {}
    return stock_list, stock_dict
def fetch_stock_data(symbol, start_date="2010-01-01", end_date=None, qfqday=True):
    """
    使用akshare库获取股票数据
    """
    if end_date is None:
        end_date = datetime.now().strftime("%Y-%m-%d")
    
    print(f"开始获取 {symbol} 数据，时间范围: {start_date} 至 {end_date}")
    
    try:
        # 转换日期格式为akshare需要的格式
        start_date_fmt = start_date.replace('-', '')
        end_date_fmt = end_date.replace('-', '')
        
        
        # 股票数据
        ak_symbol = symbol
        print(f"获取股票数据2: {ak_symbol}")
        # 使用akshare获取股票数据
        data = ak.stock_zh_a_hist(symbol=ak_symbol, period="daily", 
                                start_date=start_date_fmt, end_date=end_date_fmt) 
                                # adjust="none")
        
        
        bar_data = []
        if not data.empty:
            print(f"原始数据条数: {len(data)}")
            # 转换列名并处理数据
            for _, row in data.iterrows():
                # 获取成交额，如果不存在则填充-1
                amount = float(row['成交额']) if '成交额' in row and pd.notna(row['成交额']) else -1
                # 获取换手率，如果不存在则填充-1
                turnover = float(row['换手率']) if '换手率' in row and pd.notna(row['换手率']) else -1
                
                # 对价格数据进行处理，如果小于0则默认为0.01
                open_price = float(row['开盘'])
                high_price = float(row['最高'])
                low_price = float(row['最低'])
                close_price = float(row['收盘'])
                
                # 确保价格不小于0.01
                open_price = max(open_price, 0.01)
                high_price = max(high_price, 0.01)
                low_price = max(low_price, 0.01)
                close_price = max(close_price, 0.01)
                
                bar_data.append({
                    'Date': pd.to_datetime(row['日期']),
                    'Open': open_price,
                    'High': high_price,
                    'Low': low_price,
                    'Close': close_price,
                    'Adj Close': close_price,  # akshare的后复权数据已经处理过
                    'Volume': int(row['成交量']),
                    'Amount': amount,
                    'Turnover': turnover
                })
        else:
            print("未获取到数据")
        
        print(f"共获取到 {len(bar_data)} 条数据")
        return bar_data
    
    except Exception as e:
        print(f"获取数据出错: {e}")
        return []



def save_to_csv(data, symbol):
    """将股票数据保存为CSV文件，支持追加数据"""
    if not data:
        print("没有数据需要保存为CSV文件")
        return
    
    # 创建DataFrame
    df = pd.DataFrame(data)
    
    # 将datetime列格式化为字符串
    df['Date'] = df['Date'].dt.strftime('%Y-%m-%d')
    # 修复Volume列转换错误：先转换为浮点数，再转换为整数
    df['Volume'] = df['Volume'].astype(float).astype(int)
    # 确保Amount和Turnover列存在，如果不存在则添加并填充-1
    if 'Amount' not in df.columns:
        df['Amount'] = -1
    if 'Turnover' not in df.columns:
        df['Turnover'] = -1
    # 创建保存目录（如果不存在）
    csv_dir = 'csv_data'
    if not os.path.exists(csv_dir):
        os.makedirs(csv_dir)
        print(f"创建目录: {csv_dir}")
    
    # 构建文件名
    file_name = f"{csv_dir}/{symbol}_stock_data.csv"
    
    # 检查文件是否存在
    if os.path.exists(file_name):
        print(f"CSV文件已存在: {file_name}，将尝试追加数据")
        try:
            # 读取已有数据
            existing_df = pd.read_csv(file_name, encoding='utf-8-sig')
            # 合并新数据和已有数据
            combined_df = pd.concat([existing_df, df])
            # 根据Date列去重，保留最后一条记录（如果有重复）
            combined_df.drop_duplicates(subset=['Date'], keep='last', inplace=True)
            # 按日期排序
            combined_df.sort_values('Date', inplace=True)
            # 保存合并后的数据
            combined_df.to_csv(file_name, index=False, encoding='utf-8-sig')
            print(f"数据已追加到CSV文件: {file_name}")
        except Exception as e:
            print(f"追加数据出错: {e}")
            # 出错时仍然尝试保存新数据（覆盖模式）
            df.to_csv(file_name, index=False, encoding='utf-8-sig')
            print(f"数据已保存到CSV文件（覆盖模式）: {file_name}")
    else:
        # 文件不存在，创建新文件
        df.to_csv(file_name, index=False, encoding='utf-8-sig')
        print(f"数据已保存到CSV文件: {file_name}")

def get_last_record_date(symbol):
    """获取CSV文件中某只股票的最后记录日期"""
    # 构建CSV文件路径
    csv_dir = 'csv_data'
    file_name = f"{csv_dir}/{symbol}_stock_data.csv"
    
    # 检查文件是否存在
    if not os.path.exists(file_name):
        print(f"CSV文件不存在: {file_name}")
        return None
    
    try:
        # 使用pandas读取CSV文件
        df = pd.read_csv(file_name, encoding='utf-8-sig')
        
        # 如果文件为空或只有表头
        if len(df) == 0:
            print(f"CSV文件为空: {file_name}")
            return None
        
        # 获取最后一行的Date列值
        last_record_date_str = df['Date'].iloc[-1]
        print(f"从CSV文件读取的最后记录日期: {last_record_date_str}")
        
        # 转换为datetime对象并返回
        last_record_date = datetime.strptime(last_record_date_str, "%Y-%m-%d")
        return last_record_date
    except Exception as e:
        print(f"读取CSV文件最后记录日期出错: {e}")
        return None

# def is_trade_day_for_date(date_obj):
#     """
#     检查指定日期是否为交易日（使用与is_trade_day相同的逻辑）
#     """
#     # weekday() 返回 0-6 (周一到周日), 所以 +1 后是 1-7 (周一到周日)
#     if date_obj.weekday()+1 in [1, 3, 5]:  # 修改为周一、周三、周五
#         # 简单判断，周一到周五为工作日（不考虑节假日）
#         if date_obj.weekday() < 5:  # 0-4 代表周一到周五
#             return True
#         else:
#             return False
#     else:
#         return False

def has_trade_days_between(start_date_str):
    """
    检查从start_date到当前日期之间是否存在交易日
    """
    try:
        start_date = datetime.strptime(start_date_str, "%Y-%m-%d")
        current_date = datetime.now()
        # current_date = datetime.now() - timedelta(days=10)#for test

        
        # 遍历从start_date到current_date的每一天
        check_date = start_date
        while check_date <= current_date:
            if is_trade_day(check_date):
                return True
            check_date += timedelta(days=1)
        
        return False
    except Exception as e:
        print(f"检查交易日出错: {e}")
        return False

def bk_pull_stock_data(symbol):
    # 1. 获取最后记录日期，确定下载范围
    last_date = get_last_record_date(symbol)
    # 如果有最后记录日期，则从该日期的下一天开始获取数据；否则使用默认日期
    if last_date:
        start_date = (last_date + timedelta(days=1)).strftime("%Y-%m-%d")
        print(f"最后记录日期: {last_date.date()}")
        print(f"将从 {start_date} 开始获取数据")
        # 使用函数 is_trade_day 判断日期start_date到当前日期之间,是否存在交易日,如果不存在直接返回
        if not has_trade_days_between(start_date):
            print(f"从 {start_date} 到当前日期之间没有交易日，直接返回")
            return
        
    else:
        # 修改默认开始日期为当前日期减去5年
        start_date = (datetime.now() - timedelta(days=3*365)).strftime("%Y-%m-%d")
        print(f"-没有找到最后记录日期,默认日期 {start_date} 开始获取数据")

    # 2. 获取股票数据
    stock_data = fetch_stock_data(symbol, start_date=start_date)
    
    # 打印准备保存的数据信息
    print(f"准备保存的数据: {len(stock_data)} 条，数据样本: {stock_data[:1] if stock_data else '无'}")
    
    # 3. 保存到CSV文件
    save_to_csv(stock_data, symbol)

config_file = 'user_cfg.json'
def download_stock_data_usercfg():
    """使用用户配置文件下载股票数据"""
    try:
        with open(config_file, 'r') as f:
            config = json.load(f)
        # 修改为使用stock_list而不是symbols
        for symbol in config['stock_list']:
            print(f"开始下载股票数据: {symbol}")
            bk_pull_stock_data(symbol)
    except Exception as e:
        print(f"读取配置文件出错: {e}")

    pass

def download_stock_data_from_stocklist(stock_list):
    """使用传入的股票列表下载股票数据"""
    try:
        # 直接使用传入的stock_list参数进行循环处理
        for item in stock_list:
            # 处理元组格式(code, name)或简单字符串
            if isinstance(item, tuple) and len(item) >= 1:
                symbol = item[0]
                name = item[1] if len(item) > 1 else ""
            else:
                symbol = str(item)
                name = ""
            
            # 去除空白字符
            symbol = symbol.strip()
            # 确保代码不为空
            if symbol:
                # 检查是否需要去掉前两个字符（大小写不敏感）
                symbol_lower = symbol.lower()
                if len(symbol) > 2 and (symbol_lower.startswith('sh') or symbol_lower.startswith('sz')):
                    stock_code = symbol[2:]
                    if name:
                        print(f"开始下载股票数据: {stock_code} - {name} (原代码: {symbol})")
                    else:
                        print(f"开始下载股票数据: {stock_code} (原代码: {symbol})")
                    bk_pull_stock_data(stock_code)
                else:
                    if name:
                        print(f"开始下载股票数据: {symbol} - {name}")
                    else:
                        print(f"开始下载股票数据: {symbol}")
                    bk_pull_stock_data(symbol)
    except Exception as e:
        print(f"处理stock_list出错: {e}")

if __name__ == "__main__":
    # 检查是否提供了命令行参数
    if len(sys.argv) > 1:
        symbol = sys.argv[1]
        print(f"get stock date: {symbol}")
        bk_pull_stock_data(symbol)
    else:
      
        print("未提供股票代码，将使用文件中的股票列表")
        # 使用read_list_file函数读取股票列表文件
        stock_list, stock_dict = read_list_file("Table.txt")
        download_stock_data_from_stocklist(stock_list)

